from launch import LaunchDescription
from launch_ros.actions import Node
from launch.actions import DeclareLaunchArgument, OpaqueFunction
from launch.substitutions import LaunchConfiguration


def generate_nodes(context):
    """根据 uav_count 动态生成节点列表"""
    # 解析运行时参数
    uav_count = int(LaunchConfiguration('uav_count').perform(context))  # 获取 UAV 数量
    gcs_ip = LaunchConfiguration('gcs_ip').perform(context)  # 获取地面站 IP
    px4_sim_host = LaunchConfiguration('px4_sim_host').perform(context)  # 获取 PX4 仿真主机地址

    # 动态生成节点
    nodes = []
    for i in range(uav_count):
        node = Node(
            package='mavros',
            executable='mavros_node',
            namespace=f'/drone_{i}',
            output='screen',
            emulate_tty=True,
            parameters=[
                {'fcu_url': f'udp://:{14540 + i}@{px4_sim_host}:{14580 + i}'},
                {'tgt_system': i + 1},
                {'gcs_url': f'udp://@{gcs_ip}:{14550 + i}'}
            ]
        )
        nodes.append(node)
    return nodes


def generate_launch_description():
    # 声明参数
    uav_count_arg = DeclareLaunchArgument(
        'uav_count',
        default_value='3',  # 默认生成 3 个节点
        description='Number of UAV nodes to launch.'
    )

    gcs_ip_arg = DeclareLaunchArgument(
        'gcs_ip',
        default_value='192.168.16.191',
        description='Ground Control Station IP address.'
    )

    px4_sim_host_arg = DeclareLaunchArgument(
        'px4_sim_host',
        default_value='localhost',
        description='PX4 simulation host address.'
    )

    # 返回 LaunchDescription，动态生成节点通过 OpaqueFunction 实现
    return LaunchDescription([
        uav_count_arg,
        gcs_ip_arg,
        px4_sim_host_arg,
        OpaqueFunction(function=generate_nodes)  # 使用 OpaqueFunction 动态生成节点
    ])
